# -*- coding: utf-8 -*-
from functools import wraps

from loguru import logger

from db.base import SessionLocal


def with_session(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        session = SessionLocal()
        try:
            result = f(session, *args, **kwargs)
            session.commit()
            return result
        except Exception as e:
            try:
                session.rollback()
            except Exception as rb_err:
                # 忽略回滚过程中的连接层错误，避免再次抛出
                logger.error(rb_err)
            logger.error(e)
            raise
        finally:
            session.close()
    return wrapper


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()
